{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting pygraphviz\n",
      "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/7e/b1/d6d849ddaf6f11036f9980d433f383d4c13d1ebcfc3cd09bc845bda7e433/pygraphviz-1.5.zip (117kB)\n",
      "\u001b[K    100% |████████████████████████████████| 122kB 3.0MB/s ta 0:00:01\n",
      "\u001b[?25hBuilding wheels for collected packages: pygraphviz\n",
      "  Building wheel for pygraphviz (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Stored in directory: /Users/saul/Library/Caches/pip/wheels/65/54/69/1aee9e66ab19916293208d4c9de0d3898adebe6b2eeff6476b\n",
      "Successfully built pygraphviz\n",
      "Installing collected packages: pygraphviz\n",
      "Successfully installed pygraphviz-1.5\n"
     ]
    }
   ],
   "source": [
    "!pip install pygraphviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pygraphviz as pgv\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "REMOVE_NODES = {\n",
    "    \"MIOpen\",\n",
    "    \"AMD Math Library\",\n",
    "    \"ROCDL\",\n",
    "    \"rocBLAS\",\n",
    "    \"ONNX Runtime\",\n",
    "    \"clDNN\",\n",
    "    \"MKL\",\n",
    "    \"OpenVino\",\n",
    "    \"BigDL\",\n",
    "    \"nGraph\",\n",
    "    \"iDeep\",\n",
    "    \"MKL\",\n",
    "    \"clDNN\",\n",
    "    \"Metal\",\n",
    "    \"MLIR: nGraph\",\n",
    "    \"DNNL\",\n",
    "    \"PlaidML\",\n",
    "    \"NNVM\",\n",
    "    \"OpenCL\",\n",
    "    \"Vulkan\",\n",
    "    \"SPIR-V\",\n",
    "    \"OpenGL\",\n",
    "    \"VTA\",\n",
    "    \"MLIR: TensorFlow\",\n",
    "    \"MLIR: LLVM\",\n",
    "    \"MLIR: ROCDL\",\n",
    "    \"MLIR: SPIR-V\",\n",
    "    \"MLIR: Vector\",\n",
    "    \"MLIR: Standard\",\n",
    "    \"MLIR: GPU\",\n",
    "    \"MLIR: NVVM\",\n",
    "    \"MLIR: CUDA\",\n",
    "    \"MLIR: Loop\",\n",
    "    \"MLIR: Affine\",\n",
    "    \"CPU\",\n",
    "    \"ROCm\",\n",
    "    \"CUDA\",\n",
    "    \"LLVM\",\n",
    "    \"OpenCV\",\n",
    "    \"NNPACK\",\n",
    "    \"BLAS\",\n",
    "    \"Relay\",\n",
    "    \"TensorRT\",\n",
    "    \"NVVM\",\n",
    "    \"cuDNN\",\n",
    "    \"ChainerX\",\n",
    "    \"Caffe\",\n",
    "    \"BigQuant\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "G=pgv.AGraph(\"graph.dot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "for n in REMOVE_NODES:\n",
    "    in_ = G.in_neighbors(n)\n",
    "    out_ = G.out_neighbors(n)\n",
    "    for edge in itertools.product(in_, out_):\n",
    "        if edge not in G.edges():\n",
    "            G.add_edge(*edge)\n",
    "    G.delete_node(n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove all edges and re-add so they qre unique"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {},
   "outputs": [],
   "source": [
    "edges = set(G.edges())\n",
    "G.delete_edges_from(edges)\n",
    "# Remove self edges\n",
    "edges = {(l, r) for l, r in edges if l != r}\n",
    "G.add_edges_from(edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {},
   "outputs": [],
   "source": [
    "G.write(\"out.dot\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
